/*
池化层实现
实现最大池化和平均池化操作
支持特征图下采样和维度管理
*/
#include "PoolingLayer.h"
#include "ConvLayer.h"
#include <string>
#include <iostream>
#include <stdexcept>

void Pooling(const PoolingType& poolingType, const Tensor<double, 3>& inputFeatureMap, Tensor<double, 3>& outputFeatureMap, int widthPoolingSize, int heightPoolingSize, int stride) {
    int inputWidth = inputFeatureMap.dimension(0);
    int inputHeight = inputFeatureMap.dimension(1);
    int inputChannels = inputFeatureMap.dimension(2);
    int outputWidth = (inputWidth - widthPoolingSize) / stride + 1;
    int outputHeight = (inputHeight - heightPoolingSize) / stride + 1;
    int outputChannels = inputChannels;
    outputFeatureMap = Tensor<double, 3>(outputWidth, outputHeight, outputChannels);
    outputFeatureMap.setZero();
    
    for (int c = 0; c < outputChannels; ++c) {
        for (int h = 0; h < outputHeight; ++h) {
            for (int w = 0; w < outputWidth; ++w) {
                for (int i = 0; i < widthPoolingSize; ++i) {
                    for (int j = 0; j < heightPoolingSize; ++j) {
                        if (poolingType == PoolingType::MAX_POOLING) {
                            outputFeatureMap(w, h, c) = max(outputFeatureMap(w, h, c), inputFeatureMap(w * stride + i, h * stride + j, c));
                        }
                        else if (poolingType == PoolingType::MIN_POOLING) {
                            outputFeatureMap(w, h, c) = min(outputFeatureMap(w, h, c), inputFeatureMap(w * stride + i, h * stride + j, c));
                        }
                        else if (poolingType == PoolingType::AVG_POOLING) {
                            outputFeatureMap(w, h, c) += inputFeatureMap(w * stride + i, h * stride + j, c);
                        }
                    }
                }
                if (poolingType == PoolingType::AVG_POOLING) {
                    outputFeatureMap(w, h, c) /= widthPoolingSize * heightPoolingSize;
                }
            }
        }
    }
}

PoolingLayer::PoolingLayer(int layerIndex, const LayerType& layerType, int inputWidth, int inputHeight, int inputChannels, const PoolingType& poolingType, int widthPoolingSize, int heightPoolingSize, int stride)
    : Layer(layerIndex, layerType), inputWidth(inputWidth), inputHeight(inputHeight), inputChannels(inputChannels), poolingType(poolingType), widthPoolingSize(widthPoolingSize), heightPoolingSize(heightPoolingSize), stride(stride) {
    if (widthPoolingSize <= 0 || heightPoolingSize <= 0 || stride < 0) {
        throw invalid_argument("Invalid pooling layer parameters of pooling layer, " + to_string(layerIndex) + "width pollingSize: " + to_string(widthPoolingSize) + ", height pollingSize: " + to_string(heightPoolingSize) + ", stride: " + to_string(stride));
    }
    if (layerIndex <= 0 || layerIndex >= LayerManager::getInstance().getNumLayers()) {
        throw invalid_argument("Invalid layer index of pooling layer, " + to_string(layerIndex) + ", numLayers: " + to_string(LayerManager::getInstance().getNumLayers()));
    }
    calculateOutputDimensions(inputWidth, inputHeight, inputChannels, widthPoolingSize, heightPoolingSize, stride);
    outputFeatureMap = Tensor<double, 3>(outputWidth, outputHeight, outputChannels);
    outputFeatureMap.setZero();
    outputFeatureMapForConv = Tensor<double, 3>(outputWidth, outputHeight, outputChannels);
    outputFeatureMapForConv.setZero();
    outputFeatureMapForFc = Tensor<double, 1>(outputWidth * outputHeight * outputChannels);
    outputFeatureMapForFc.setZero();
    inputsGradient = Tensor<double, 3>(inputWidth, inputHeight, inputChannels);
    inputsGradient.setZero();
    outputsGradient = Tensor<double, 3>(outputWidth, outputHeight, outputChannels);
    outputsGradient.setZero();
    weightsGradient = Tensor<double, 3>(widthPoolingSize, heightPoolingSize, inputChannels, outputChannels);
    weightsGradient.setZero();
    biasesGradient = Tensor<double, 1>(outputChannels);
    biasesGradient.setZero();
    inputsDelta = Tensor<double, 3>(inputWidth, inputHeight, inputChannels);
    inputsDelta.setZero();
    outputsDelta = Tensor<double, 3>(outputWidth, outputHeight, outputChannels);
    outputsDelta.setZero();
    weightsDelta = Tensor<double, 3>(widthPoolingSize, heightPoolingSize, inputChannels, outputChannels);
    weightsDelta.setZero();
    biasesDelta = Tensor<double, 1>(outputChannels);
    biasesDelta.setZero();
    prevBatchInputsDelta = Tensor<double, 3>(inputWidth, inputHeight, inputChannels);
    prevBatchInputsDelta.setZero();
    prevBatchOutputsDelta = Tensor<double, 3>(outputWidth, outputHeight, outputChannels);
    prevBatchOutputsDelta.setZero();
    prevBatchWeightsDelta = Tensor<double, 3>(widthPoolingSize, heightPoolingSize, inputChannels, outputChannels);
    prevBatchWeightsDelta.setZero();
    prevBatchBiasesDelta = Tensor<double, 1>(outputChannels);
    prevBatchBiasesDelta.setZero();
}
PoolingLayer::~PoolingLayer() {}
void PoolingLayer::setPrevLayer(const shared_ptr<Layer>& prevLayer) {
    this->prevLayer = prevLayer;
    if (prevLayer->getLayerType() == LayerType::CONV) {
        this->prevConvLayer = dynamic_pointer_cast<ConvLayer>(prevLayer);
    } else {
        throw invalid_argument("Invalid layer type when setting previous layer of pooling layer " + to_string(layerIndex) + " : " + to_string(prevLayer->getLayerType()));
    }
}
void PoolingLayer::setNextLayer(const shared_ptr<Layer>& nextLayer) {
    this->nextLayer = nextLayer;
    if (nextLayer->getLayerType() == LayerType::CONV) {
        this->nextConvLayer = dynamic_pointer_cast<ConvLayer>(nextLayer);
    } else if (nextLayer->getLayerType() == LayerType::FC) {
        this->nextFullConnectedLayer = dynamic_pointer_cast<FullConnectedLayer>(nextLayer);
    } else {
        throw invalid_argument("Invalid layer type when setting next layer of pooling layer " + to_string(layerIndex) + " : " + to_string(prevLayer->getLayerType()));
    }
}
const PoolingType& PoolingLayer::getPoolingType() const {
    return poolingType;
}
int PoolingLayer::getWidthPoolingSize() const {
    return widthPoolingSize;
}
int PoolingLayer::getHeightPoolingSize() const {
    return heightPoolingSize;
}
int PoolingLayer::getStride() const {
    return stride;
}
const Tensor<double, 3>& PoolingLayer::getInputFeatureMap() const {
    return inputFeatureMap;
}
int PoolingLayer::getInputWidth() const {
    return inputWidth;
}
int PoolingLayer::getInputHeight() const {
    return inputHeight;
}
int PoolingLayer::getInputChannels() const {
    return inputChannels;
}
const Tensor<double, 3>& PoolingLayer::getOutputFeatureMap() const {
    return outputFeatureMap;
}
const Tensor<double, 3>& PoolingLayer::getOutputFeatureMapForConv() const {
    return outputFeatureMapForConv;
}
const Tensor<double, 1>& PoolingLayer::getOutputFeatureMapForFc() const {
    return outputFeatureMapForFc;
}
int PoolingLayer::getOutputWidth() const {
    return outputWidth;
}
int PoolingLayer::getOutputHeight() const {
    return outputHeight;
}
int PoolingLayer::getOutputChannels() const {
    return outputChannels;
}
int PoolingLayer::getOutputWidthForConv() const {
    return outputWidthForConv;
}
int PoolingLayer::getOutputHeightForConv() const {
    return outputHeightForConv;
}
int PoolingLayer::getOutputChannelsForConv() const {
    return outputChannelsForConv;
}
int PoolingLayer::getOutputSizeForFc() const {
    return outputSizeForFc;
}
const Tensor<double, 3>& PoolingLayer::getInputsGradient() const {
    return inputsGradient;
}
const Tensor<double, 3>& PoolingLayer::getOutputsGradient() const {
    return outputsGradient;
}
const Tensor<double, 3>& PoolingLayer::getWeightsGradient() const {
    return weightsGradient;
}
const Tensor<double, 1>& PoolingLayer::getBiasesGradient() const {
    return biasesGradient;
}
const Tensor<double, 3>& PoolingLayer::getInputsDelta() const {
    return inputsDelta;
}
const Tensor<double, 3>& PoolingLayer::getOutputsDelta() const {
    return outputsDelta;
}
const Tensor<double, 3>& PoolingLayer::getWeightsDelta() const {
    return weightsDelta;
}
const Tensor<double, 1>& PoolingLayer::getBiasesDelta() const {
    return biasesDelta;
}
void PoolingLayer::setPoolingType(const PoolingType& poolingType) {
    this->poolingType = poolingType;
}
void PoolingLayer::setWidthPoolingSize(int widthPoolingSize) {
    this->widthPoolingSize = widthPoolingSize;
}
void PoolingLayer::setHeightPoolingSize(int heightPoolingSize) {
    this->heightPoolingSize = heightPoolingSize;
}
void PoolingLayer::setStride(int stride) {
    this->stride = stride;
}
void PoolingLayer::setInputFeatureMap(const Tensor<double, 3>& inputFeatureMap) {
    this->inputFeatureMap = inputFeatureMap;
    // setInputWidth(inputFeatureMap.dimension(0));
    // setInputHeight(inputFeatureMap.dimension(1));
    // setInputChannels(inputFeatureMap.dimension(2));
}
void PoolingLayer::setInputWidth(int inputWidth) {
    this->inputWidth = inputWidth;
}
void PoolingLayer::setInputHeight(int inputHeight) {
    this->inputHeight = inputHeight;
}
void PoolingLayer::setInputChannels(int inputChannels) {
    this->inputChannels = inputChannels;
}
void PoolingLayer::setOutputFeatureMap(const Tensor<double, 3>& outputFeatureMap) {
    this->outputFeatureMap = outputFeatureMap;
}
void PoolingLayer::setOutputFeatureMapForConv(const Tensor<double, 3>& outputFeatureMapForConv) {
    this->outputFeatureMapForConv = outputFeatureMapForConv;
}
void PoolingLayer::setOutputFeatureMapForFc(const Tensor<double, 1>& outputFeatureMapForFc) {
    this->outputFeatureMapForFc = outputFeatureMapForFc;
}
void PoolingLayer::setOutputWidth(int outputWidth) {
    this->outputWidth = outputWidth;
}
void PoolingLayer::setOutputHeight(int outputHeight) {
    this->outputHeight = outputHeight;
}
void PoolingLayer::setOutputChannels(int outputChannels) {
    this->outputChannels = outputChannels;
}
void PoolingLayer::setOutputWidthForConv(int outputWidthForConv) {
    this->outputWidthForConv = outputWidthForConv;
}
void PoolingLayer::setOutputHeightForConv(int outputHeightForConv) {
    this->outputHeightForConv = outputHeightForConv;
}
void PoolingLayer::setOutputChannelsForConv(int outputChannelsForConv) {
    this->outputChannelsForConv = outputChannelsForConv;
}
void PoolingLayer::setOutputSizeForFc(int outputSizeForFc) {
    this->outputSizeForFc = outputSizeForFc;
}
void PoolingLayer::calculateOutputDimensions(int inputWidth, int inputHeight, int inputChannels, int widthPoolingSize, int heightPoolingSize, int stride) {
    this->outputWidth = (inputWidth - widthPoolingSize) / stride + 1;
    this->outputHeight = (inputHeight - heightPoolingSize) / stride + 1;
    this->outputChannels = inputChannels;
    this->outputWidthForConv = this->outputWidth;
    this->outputHeightForConv = this->outputHeight;
    this->outputChannelsForConv = this->outputChannels;
    this->outputSizeForFc = this->outputWidth * this->outputHeight * this->outputChannels;
}
void PoolingLayer::calculateOutputFeatureMap() {
    PoolingType poolingType = getPoolingType();
    int widthPoolingSize = getWidthPoolingSize();
    int heightPoolingSize = getHeightPoolingSize();
    int stride = getStride();
    Pooling(poolingType, this->inputFeatureMap, this->outputFeatureMap, widthPoolingSize, heightPoolingSize, stride);
    if (this->nextLayer->getLayerType() == LayerType::CONV) {
        transformOutputFeatureMapToConv(this->outputFeatureMap);
    } else if (this->nextLayer->getLayerType() == LayerType::FC) {
        transformOutputFeatureMapToFc(this->outputFeatureMap);
    }
}
void PoolingLayer::transformOutputFeatureMapToConv(Tensor<double, 3>& outputFeatureMap) {
    this->outputFeatureMapForConv = outputFeatureMap.reshape(outputFeatureMap.dimension(0), outputFeatureMap.dimension(1), outputFeatureMap.dimension(2));
}
void PoolingLayer::transformOutputFeatureMapToFc(Tensor<double, 3>& outputFeatureMap) {
    int totalSize = outputFeatureMap.dimension(0) * outputFeatureMap.dimension(1) * outputFeatureMap.dimension(2);
    this->outputFeatureMapForFc = outputFeatureMap.reshape(totalSize);
}
void PoolingLayer::calculateGradient(Tensor<double, 3>& calcGradient, const Tensor<double, 3>& nextGradient, const Tensor<double, 3>& featureMap, const int width, const int height, const int channels) {
    calcGradient.setZero();
    for (int c = 0; c < channels; ++c) {
        for (int h = 0; h < height; ++h) {
            for (int w = 0; w < width; ++w) {
                if (getPoolingType() == PoolingType::MAX_POOLING) {
                    double maxVal = -std::numeric_limits<double>::max();
                    int maxI = -1, maxJ = -1;
                    for (int i = 0; i < widthPoolingSize; ++i) {
                        for (int j = 0; j < heightPoolingSize; ++j) {
                            int x = w + i;
                            int y = h + j;
                            if (x >= featureMap.dimension(0) || y >= featureMap.dimension(1)) {
                                continue;
                            }
                            if (featureMap(x, y, c) > maxVal) {
                                maxVal = featureMap(x, y, c);
                                maxI = i;
                                maxJ = j;
                            }
                        }
                    }
                    if (maxI != -1 && maxJ != -1) {
                        calcGradient(w + maxI, h + maxJ, c) = nextGradient(c, h, w);
                    }
                }
                if (getPoolingType() == PoolingType::MIN_POOLING) {
                    double minVal = std::numeric_limits<double>::min();
                    int minI = -1, minJ = -1;
                    for (int i = 0; i < widthPoolingSize; ++i) {
                        for (int j = 0; j < heightPoolingSize; ++j) {
                            int x = w + i;
                            int y = h + j;
                            if (x >= featureMap.dimension(0) || y >= featureMap.dimension(1)) {
                                continue;
                            }
                            if (featureMap(x, y, c) < minVal) {
                                minVal = featureMap(x, y, c);
                                minI = i;
                                minJ = j;
                            }
                        }
                    }
                    if (minI != -1 && minJ != -1) {
                        calcGradient(w + minI, h + minJ, c) = nextGradient(c, h, w);
                    }
                }
                if (getPoolingType() == PoolingType::AVG_POOLING) {
                    double avgGradient = nextGradient(c, h, w) / (widthPoolingSize * heightPoolingSize);
                    for (int i = 0; i < widthPoolingSize; ++i) {
                        for (int j = 0; j < heightPoolingSize; ++j) {
                            int x = w + i;
                            int y = h + j;
                            if (x >= featureMap.dimension(0) || y >= featureMap.dimension(1)) {
                                continue;
                            }
                            calcGradient(x, y, c) = avgGradient;
                        }
                    }
                }
            }
        }
    }
}
void PoolingLayer::calculateInputsGradient(const Tensor<double, 3>& outputsGradient) {
    calculateGradient(this->inputsGradient, outputsGradient, this->inputFeatureMap, this->inputWidth, this->inputHeight, this->inputChannels);
}
void PoolingLayer::calculateOutputsGradient(const Tensor<double, 3>& nextLayerOutputsGradient) {
    calculateGradient(this->outputsGradient, nextLayerOutputsGradient, this->outputFeatureMap, this->outputWidth, this->outputHeight, this->outputChannels);
}
void PoolingLayer::calculateInputsDelta(const Tensor<double, 3>& outputDelta) {
    inputsDelta = outputsDelta;
    prevBatchInputsDelta = inputsDelta;
}
void PoolingLayer::calculateOutputsDelta(int learningRate, int momentum) {
    outputsDelta = outputsGradient * learningRate + prevBatchOutputsDelta * momentum;
    prevBatchOutputsDelta = outputsDelta;
}
void PoolingLayer::forward(const Tensor<double, 3>& inputFeatureMap) {
    setInputFeatureMap(inputFeatureMap);
    calculateOutputFeatureMap();
}
void PoolingLayer::backward(const Tensor<double, 3>& nextLayerOutputsGradient) {
    if (prevLayer == nullptr) {
        return;
    }
    if (prevConvLayer != nullptr) {
        if (prevConvLayer->getLayerIndex() != layerIndex - 1) {
            throw runtime_error("Previous convolution layer index is not correct.");
        }
    }
    if (nextConvLayer != nullptr) {
        if (nextConvLayer->getLayerIndex() != layerIndex + 1) {
            throw runtime_error("Next convolution layer index is not correct.");
        }
    }
    if (nextFullConnectedLayer != nullptr) {
        if (nextFullConnectedLayer->getLayerIndex() != layerIndex + 1) {
            throw runtime_error("Next full connected layer index is not correct.");
        }
    }
    calculateOutputsGradient(nextLayerOutputsGradient);
    calculateInputsGradient(this->getOutputsGradient());
}
void PoolingLayer::print() const {
    string indent = "    ";
    cout << "Pooling layer " << layerIndex << ":" << endl;
    cout << "Input size: " << inputWidth << "x" << inputHeight << "x" << inputChannels << endl;
    cout << "Output size: " << outputWidth << "x" << outputHeight << "x" << outputChannels << endl;
    cout << "Pooling type: " << poolingType << endl;
    cout << "Pooling size: " << widthPoolingSize << "x" << heightPoolingSize << endl;
    cout << "Stride: " << stride << endl;
    cout << "Next layer type: " << to_string(nextLayerType) << endl;
    cout << "Gradients: " << endl;
    cout << "Inputs gradient: " << endl;
    for (int c = 0; c < inputChannels; ++c) {
        for (int h = 0; h < inputHeight; ++h) {
            for (int w = 0; w < inputWidth; ++w) {
                cout << indent << inputsGradient(c, h, w) << " ";
            }
            cout << endl;
        }
        cout << endl;
    }
    cout << "Outputs gradient: " << endl;
    for (int c = 0; c < outputChannels; ++c) {
        for (int h = 0; h < outputHeight; ++h) {
            for (int w = 0; w < outputWidth; ++w) {
                cout << indent << outputsGradient(c, h, w) << " ";
            }
            cout << endl;
        }
        cout << endl;
    }
    cout << "Deltas: " << endl;
    cout << "Inputs delta: " << endl;
    for (int c = 0; c < inputChannels; ++c) {
        for (int h = 0; h < inputHeight; ++h) {
            for (int w = 0; w < inputWidth; ++w) {
                cout << indent << inputsDelta(c, h, w) << " ";
            }
            cout << endl;
        }
        cout << endl;
    }
    cout << "Outputs delta: " << endl;
    for (int c = 0; c < outputChannels; ++c) {
        for (int h = 0; h < outputHeight; ++h) {
            for (int w = 0; w < outputWidth; ++w) {
                cout << indent << outputsDelta(c, h, w) << " ";
            }
            cout << endl;
        }
        cout << endl;
    }
}